# -*- codeing = utf-8 -*-
# @Time : 2022/2/18 13:50
# @File : recognize_face.py
# @Software : PyCharm
import cv2
import time
from train import predict_model
import numpy as np
from PIL import Image, ImageFont, ImageDraw
import database

FACE_LABEL = database.dataoperator().get_recognize_info()
# print(FACE_LABEL)
C_flag = np.zeros(database.dataoperator().getpersoncount(), dtype=np.int32)


def recognize_video(ui, window_name='face recognize', camera_idx=0):
    cv2.namedWindow(window_name)

    cap = cv2.VideoCapture(camera_idx)
    count = 0
    while cap.isOpened():
        ok, frame = cap.read()
        if not ok:
            break

        flag, catch_frame, label = catch_face(frame)
        if flag == True:
            if count == 0:
                start = time.time()
            count += 1
        else:
            count = 0
        c = cv2.waitKey(1)
        cv2.imshow(window_name, catch_frame)
        if count == 3:
            end = time.time()
            count = 0
            C_flag[label] = 1
            ui.printf("打卡成功 , 花费时常{:.3f} m".format(end - start))

        if c & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()


def catch_face(frame):
    classfier = cv2.CascadeClassifier(
        "D:\\Anaconda\\Lib\\site-packages\\cv2\\data\\haarcascade_frontalface_alt2.xml")

    grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    face_rects = classfier.detectMultiScale(grey, scaleFactor=1.2, minNeighbors=3, minSize=(32, 32))
    flag = False
    label = 0
    if len(face_rects) > 0:
        for face_rects in face_rects:
            x, y, w, h = face_rects
            image = frame[y - 10:y + h + 10, x - 10:x + w + 10]
            PIL_image = cv2pil(image)
            label = predict_model(PIL_image)
            # print(label)
            if label == 0:
                flag = False
                name = FACE_LABEL[0]
                color = (255, 0, 0)
            else:
                if C_flag[label] == 1:
                    flag = False
                    name = "请勿重复打卡"
                    color = (0, 255, 255)
                else:
                    flag = True
                    name = FACE_LABEL[label]
                    color = (0, 255, 0)

            cv2.rectangle(frame, (x - 10, y - 10), (x + w + 10, y + h + 10), color, 2)
            frame = paint_chinese_opencv(frame, name, (x - 10, y + h + 10), color)
            # ui.printf("检测到 " + str(label) + " 号 ")

        cv2.imwrite("data/tmp/{}.jpg".format(int(time.time())), frame)
    return flag, frame, label


# 摄像头使用的Opencv，pytorch识别图像使用PIL模式，需要转换
def cv2pil(image):
    return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))


# 给预测的人，打赏中文标
def paint_chinese_opencv(im, chinese, pos, color):
    img_PIL = Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
    font = ImageFont.truetype('font_s.ttf', 20)
    fillColor = color
    position = pos
    if not isinstance(chinese, str):
        chinese = chinese.decode('utf-8')
    draw = ImageDraw.Draw(img_PIL)
    draw.text(position, chinese, font=font, fill=fillColor)

    img = cv2.cvtColor(np.asarray(img_PIL), cv2.COLOR_RGB2BGR)
    return img
